package com.idanchuang.component.consumer;

import com.alibaba.fastjson.JSON;
import com.idanchuang.component.consumer.config.RpcConfig;
import com.idanchuang.component.consumer.servlet.ConsumerServletRequestContext;
import com.idanchuang.component.core.Dc;
import com.idanchuang.component.core.ServiceInfo;
import com.idanchuang.component.core.helper.EnvContext;
import com.idanchuang.component.logback.helper.LogContext;
import com.idanchuang.component.provider.api.util.RpcThreadLocalUtil;
import com.idanchuang.component.provider.api.util.RpcUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.http.client.support.HttpRequestWrapper;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.IOException;
import java.net.URI;
import java.net.URLEncoder;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author yjy
 * Created at 2021/7/19 2:18 下午
 */
public class RestTemplateInterceptor implements ClientHttpRequestInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(RestTemplateInterceptor.class);

    private final RpcConfig rpcConfig;
    private final ConsumerServletRequestContext consumerServletRequestContext;

    public RestTemplateInterceptor(RpcConfig rpcConfig, ConsumerServletRequestContext consumerServletRequestContext) {
        this.rpcConfig = rpcConfig;
        this.consumerServletRequestContext = consumerServletRequestContext;
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
        try {
            String path = request.getURI().getPath();
            //这里可以添加feign请求的全局参数
            String secretKey = rpcConfig.getSecretKey();
            List<String> extraSecretKeys = rpcConfig.getExtraSecretKeys();
            String[] signs = new String[1 + extraSecretKeys.size()];
            signs[0] = RpcUtil.sign(path, secretKey);
            for (int i = 0; i < extraSecretKeys.size(); i++) {
                signs[i + 1] = RpcUtil.sign(path, extraSecretKeys.get(i));
            }
            request.getHeaders().addAll(rpcConfig.getSignKey(), Arrays.asList(signs));
            // 扩展请求头
            if (!rpcConfig.getHeaders().isEmpty()) {
                for (Map.Entry<String, String> kv : rpcConfig.getHeaders().entrySet()) {
                    request.getHeaders().add(kv.getKey(), encodeUrl(kv.getValue()));
                }
            }
            // 将本地线程变量表存入header传递
            Map<String, String> map = RpcThreadLocalUtil.getAllFromThreadLocal();
            if (!map.isEmpty()) {
                String mapString = JSON.toJSONString(map);
                request.getHeaders().add(rpcConfig.getThreadLocalHeaderName(), encodeUrl(mapString));
            }
            // 添加来源参数
            request.getHeaders().add(Dc.RPC_FROM_NAME, ServiceInfo.getAppId());
            String currentPath = EnvContext.getCurrentPath();
            if (currentPath != null) {
                request.getHeaders().add(Dc.RPC_FROM_API_NAME, currentPath);
            }
            request.getHeaders().add(Dc.RPC_FROM_GROUP_NAME, ServiceInfo.getGroup());
            request.getHeaders().add(Dc.RPC_FROM_VERSION_NAME, ServiceInfo.getVersion());
            // query中放一个rpc_from
            UriComponentsBuilder uriBuilder = UriComponentsBuilder.fromUri(request.getURI());
            uriBuilder.replaceQueryParam(Dc.RPC_FROM_NAME, encodeUrl(ServiceInfo.getAppId()));
            // 扩展请求参数
            if (!rpcConfig.getQueries().isEmpty()) {
                for (Map.Entry<String, Collection<String>> kv : rpcConfig.getQueries().entrySet()) {
                    Collection<String> vs = kv.getValue();
                    vs = vs.stream().map(RestTemplateInterceptor::encodeUrl).collect(Collectors.toList());
                    uriBuilder.replaceQueryParam(kv.getKey(), vs);
                }
            }
            request = new UriModifyHttpRequestWrapper(request, uriBuilder.build(true).toUri());

            // 如果存在request上下游
            if (consumerServletRequestContext != null) {
                Set<String> ths = rpcConfig.getAllTransmitHeaders();
                if (!ths.isEmpty()) {
                    // 从当前请求上下文中获取header传递至下游
                    Map<String, List<String>> originHeaders = consumerServletRequestContext.getMultiHeaders(ths);
                    Set<String> currentHeaders = request.getHeaders().keySet();
                    final HttpHeaders httpHeaders = request.getHeaders();
                    originHeaders.forEach((k, v) -> {
                        if (v != null && !v.isEmpty()) {
                            if (!currentHeaders.contains(k)) {
                                httpHeaders.addAll(k, v);
                            } else {
                                logger.debug("transmit header canceled, cause currentHeaders contains key:{}", k);
                            }
                        }
                    });
                }
                Set<String> tqs = rpcConfig.getAllTransmitQueries();
                if (!tqs.isEmpty()) {
                    // 从当前请求上下文中获取queries传递至下游
                    Map<String, String[]> originQueries = consumerServletRequestContext.getMultiQueries(tqs);
                    if (!originQueries.isEmpty()) {
                        UriComponentsBuilder uriBuilder1 = UriComponentsBuilder.fromUri(request.getURI());
                        originQueries.forEach((k, v) -> {
                            if (v != null && v.length > 0) {
                                for (int i = 0; i < v.length; i++) {
                                    v[i] = encodeUrl(v[i]);
                                }
                                uriBuilder1.replaceQueryParam(k, v);
                            }
                        });
                        request = new UriModifyHttpRequestWrapper(request, uriBuilder1.build(true).toUri());
                    }
                }
            }
        } catch (Exception e) {
            // 降级
            LogContext.setKeyword("componentError");
            logger.error("RestTemplate Interceptor error", e);
        }
        return execution.execute(request, body);
    }

    private static String encodeUrl(String origin) {
        try {
            return URLEncoder.encode(origin, "UTF-8");
        } catch (Exception e) {
            logger.error("encodeUrl failed, origin: {}", origin, e);
        }
        return origin;
    }

    //...
    private static class UriModifyHttpRequestWrapper extends HttpRequestWrapper {
        private final URI uri;

        public UriModifyHttpRequestWrapper(HttpRequest request, URI uri) {
            super(request);
            this.uri = uri;
        }

        @Override
        public URI getURI() {
            return uri;
        }
    }

}
